from ase.io import read, write
import numpy as np
from pathlib import Path
import os
"""
读取OUTCAR文件：从指定目录如（E:/Li2TiO3/OUTCAR）中查找所有以"OUTCAR_"开头的文件
提取结构数据：使用ASE库读取这些OUTCAR文件中的原子结构信息（帧）
数据集分割：将所有读取到的结构数据随机打乱后，按80:20的比例分割成训练集和测试集
保存为XYZ格式：将分割后的数据集保存为XYZ格式文件，便于后续的机器学习训练使用
"""

# 读取所有 OUTCAR 文件，划分成训练集和测试集
def split_datas(root_dir,train_ratio=0.8):
    # output_dir = '../mace_data_init'
    # os.makedirs(output_dir, exist_ok=True)
    output_dir = Path('../mace_data_init')
    output_dir.mkdir(parents=True, exist_ok=True)
    outcar_files = list(root_dir.glob('OUTCAR_*'))
    frames = []
    for file in outcar_files:
        try:
            frames.extend(read(str(file), index=':'))
        except Exception as e:
            print(f"跳过文件 {file}，错误：{e}")

    # 检查是否有成功读取的帧
    if not frames:
        print("没有成功读取任何帧，检查 OUTCAR 文件的完整性和路径。")
        exit()

    # 随机打乱帧
    np.random.shuffle(frames)

    # 计算分割点
    split_index = int(train_ratio * len(frames))

    # 分割为训练集和测试集
    train_frames = frames[:split_index]
    test_frames = frames[split_index:]

    # 写入文件
    write(output_dir/ 'train.xyz', train_frames, format='extxyz')
    write(output_dir/ 'test.xyz', test_frames, format='extxyz')

    print(f'训练集包含 {len(train_frames)} 帧，测试集包含 {len(test_frames)} 帧。')

if __name__ == '__main__':
    root_dir = Path('E:/Li2TiO3/OUTCAR')
    train_ratio = 0.8
    split_datas(root_dir,train_ratio)